#
#  mfuncs.py
#  
#
#  Created by Simon Galbraith on 5/6/08.
#  Copyright (c) 2008 yCompanyName__. All rights reserved.
#
from Matrix import Matrix;
import cern,java
from cern.colt.matrix import DoubleMatrix2D as TMatrix;
from cern.colt.matrix.DoubleFactory2D import *;

try:
	from cern.colt.matrix.linalg import DoubleAlgebra as Algebra;
	from cern.jet.stat.DoubleDescriptive import sum as vsum
	import cern.jet.math.DoubleFunctions as Functions
	import cern.jet.math.DoubleMult as Mult
	import cern.jet.stat.DoubleDescriptive as Descriptive
	
except:
	from cern.colt.matrix.linalg import Algebra;
	import cern.jet.stat.Descriptive as Descriptive
	from cern.jet.stat.Descriptive import sum as vsum
	import cern.jet.math.Functions as Functions
	import cern.jet.math.Mult as Mult
	import cern.jet.stat.Descriptive as Descriptive
	
from cern.colt.list import DoubleArrayList;
from org.python.core.exceptions import ValueError as PyValueException;
from cern.colt.matrix.doublealgo import Statistic;
from java.lang.Math import sqrt



def normp(A,P):
	[tfs,exps] = size(P)
	for j in range(0,tfs):
		##v = max(DoubleArrayList(abs(P[j,:]).getElements()))
		v = norm(P[j,:],2)
		#print "norm constant is ", v
		f = Mult
		P.viewRow(j).assign(f.div(v))
		A.viewRow(j).assign(f.mult(v))
	return (A,P)

def sign(A):
	r = A.get(0,0)
	if r < 0:
		return -1
	else:
		return 1

def stddev(A):
	
	try:
		data=DoubleArrayList(A.getElements())
	except:
		data = A.toArray()
		try:
			data=data[0]
		except:
			pass
		data=DoubleArrayList(data)
	
	sumD = vsum(data)
	ss = Descriptive.sumOfSquares(data)
	sz = len(data.elements())
	#print sumD, ss, sz	return M



def mean(A):
#	print A, dir(A)
	try:
		M=Descriptive.mean(DoubleArrayList(A.getElements()))
	except:
		m = A.toArray()
		try:
			m=m[0]
		except:
			pass
		M=Descriptive.mean(DoubleArrayList(m))
	return M

def sum(A):
	print "in sum"
	q = A.zSum()
	print "after sum"
	return q

def warnings(enable=None):
	# a general error 
	# function to contain console specific errors.. 
	pass;

def nonzeros(A):
	r=cern.colt.list.IntArrayList()
	c=cern.colt.list.IntArrayList()
	v=cern.colt.list.DoubleArrayList()
	A.getNonZeros(r,c,v)
	r = r.elements()
	c = c.elements()
	rows=False
	if c[0]==c[1]:
		rows=True
		
	idx = zip(r,c,v.elements())
	idx = filter(lambda x: x[2] != 0, idx)
	idx.sort()
	if rows:
		idx = map(lambda x: x[0], idx)
	else:
		idx = map(lambda x: x[1], idx)
		
	#print "found nonzeros", idx
	return idx;
		

def norm(A,ntype=None):
	F = Algebra(0.0);
	[rows,cols] = size(A)
	r=None
	if ntype=='fro':
		r=F.normF(A);
	if ntype == 2 or ntype is None:
		C = A.toArray()
		C=C[0]
		B = cern.colt.matrix.impl.DenseDoubleMatrix1D(C)
		r = F.norm2(B)
	return r;

def rank(A):
	
	if isinstance(A,Matrix):
		F = Algebra(0.0);
		r=F.rank(A);
		return int(r);
	else:
		raise PyValueException, "Rank function can only be called on matrix objects"

def cond(A):
	F = Algebra(0.0);
	return F.cond(A);

def size(A):
	r = A.rows();
	c = A.columns();
	return (r,c);

def transpose(A):
	F = Algebra(0.0);
	if isinstance(A,float):
		return A;
	else:
		return Matrix(F.transpose(A));

def inverse(A):
	F = Algebra(0.0);
	x=F.inverse(A);
	return Matrix(x)

def zeros(m,n):
	if isinstance(m,int) and isinstance(n,int):
		return Matrix(m,n);
	else:
		raise PyValueException, "Zeros takes m=rows and n=cols as arguments.";
	return None;

def ones(m,n):
	if isinstance(m,int) and isinstance(n,int):
		return Matrix(m,n,1);
	else:
		raise PyValueException, "The function 'Zeros' takes m=rows and n=cols as arguments.";
	return None

def eig(A):
	# check that is square
	try:
		E = cern.colt.linalg.EigenvalueDecomposition(A);
		U = E.getV();
		eigs = E.getD();
	except PyValueException, e:
		print e;
		raise PyValueException,"Error in eig, check warnings()";  
	return [Matrix(eigs),Matrix(U)];

def solve(A,B):
	F = Algebra(0.0)
	if isinstance(A,cern.colt.matrix.DoubleMatrix1D):
		try:
			c = len(A.getElements())
			r = 1;
			A = A.reshape(r,c)
		except:
			c = len(A.toArray())
			A = cern.colt.matrix.impl.DenseDoubleMatrix2D([A.toArray()])
				
		#print "A was 1D\n"
	if isinstance(B,cern.colt.matrix.DoubleMatrix1D):
		try:
			c = len(B.getElements())
			r = 1;#A.columns()
			B = B.reshape(r,c)
		except:
			c = len(B.toArray())
			B = cern.colt.matrix.impl.DenseDoubleMatrix2D([B.toArray()])
		#print "B was 1D\n"
	r=None
	try:
		QR = cern.colt.matrix.linalg.QRDecomposition(A)
		r=Matrix(QR.solve(B))
	except:
		r=Matrix(F.solve(A,B))
	return r

def solveTranspose(A,B):
	F = Algebra(0.0);
	
	if isinstance(A,cern.colt.matrix.DoubleMatrix1D):
		try:
			c = len(A.getElements())
			r = 1;
			A = A.reshape(r,c)
		except:
			c = len(A.toArray())
			A = cern.colt.matrix.impl.DenseDoubleMatrix2D([A.toArray()])
				
		#print "A was 1D\n"
	if isinstance(B,cern.colt.matrix.DoubleMatrix1D):
		try:
			c = len(B.getElements())
			r = 1;#A.columns()
			B = B.reshape(r,c)
		except:
			c = len(B.toArray())
			B = cern.colt.matrix.impl.DenseDoubleMatrix2D([B.toArray()])
		#print "B was 1D\n"
	return Matrix(F.solveTranspose(B,A));

def solveLR(A,B):
	T = A.copy()
	F = cern.colt.matrix.linalg.LUDecomposition(T);
	if isinstance(A,Matrix) and isinstance(B,float):
		return F.solve(B);
	elif isinstance(A,Matrix) and isinstance(B,Matrix):
		C = F.solve(B);
	return Matrix(C)

def cov(A):
	return Matrix(Statistic.covariance(A));

def cor(A):  # handle multidimensional matrix case
	B = cov(A);
	return  Matrix(Statistic.correlation(B));


def abs(A):
	F = Functions.abs;
	if isinstance(A,float):
		return java.lang.Math.abs(A)
	else:
		A.assign(F)
	return A;

def svd(A):
	X = cern.colt.matrix.linalg.SingularValueDecomposition(A)
	u = X.getU()
	v = X.getV()
	e = X.getS()
	return [Matrix(u),Matrix(e),Matrix(v)];
